{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Training LLMs in ANY Environment with OpenEnv\n",
        "\n",
        "## \ud83c\udfaf The Vision\n",
        "\n",
        "Imagine training language models in:\n",
        "- \ud83c\udfb0 **Card games** (BlackJack, Poker, Uno)\n",
        "- \u265f\ufe0f **Board games** (Chess, Go, Connect Four)\n",
        "- \ud83d\udcc8 **Trading simulations** (realistic market environments)\n",
        "- \ud83c\udfae **Atari games** (Pong, Breakout, Space Invaders)\n",
        "- \ud83d\udcbb **Code execution environments** (interactive debugging)\n",
        "- \ud83e\udd16 **Robotics simulations** (MuJoCo, PyBullet)\n",
        "\n",
        "---\n",
        "\n",
        "### The Problem\n",
        "\n",
        "Every RL environment has different APIs:\n",
        "- \u274c OpenSpiel uses C++ bindings\n",
        "- \u274c Atari needs ALE (Arcade Learning Environment)\n",
        "- \u274c Trading sims have custom interfaces\n",
        "- \u274c Each requires different dependencies, versions, OS compatibility\n",
        "- \u274c No isolation \u2192 crashes can corrupt your system\n",
        "\n",
        "**You spend more time wrestling with environments than training models.**\n",
        "\n",
        "---\n",
        "\n",
        "### The Solution: OpenEnv - A Universal Spec\n",
        "\n",
        "<div style='background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 25px; border-radius: 10px; color: white; margin: 20px 0;'>\n",
        "    <h3 style='margin-top: 0;'>\ud83d\ude80 OpenEnv = Universal RL Environment Interface</h3>\n",
        "    <p style='font-size: 18px; line-height: 1.8;'>\n",
        "        <b>OpenEnv is not a game engine.</b><br>\n",
        "        It's a <b>specification</b> that wraps ANY RL environment with a clean, unified API.\n",
        "    </p>\n",
        "    <ul style='font-size: 16px; line-height: 1.8;'>\n",
        "        <li><b>70+ environments</b> (OpenSpiel, Atari, FinRL, and more)</li>\n",
        "        <li><b>Unified Simplified API:</b> <code>reset()</code>, <code>step(action)</code>, <code>state()</code></li>\n",
        "        <li><b>HTTP-based</b> \u2192 language-agnostic (Python, Rust, JavaScript, anything)</li>\n",
        "        <li><b>Docker-isolated</b> \u2192 reproducible, secure, no dependency hell</li>\n",
        "    </ul>\n",
        "    <p style='font-size: 16px; margin-top: 15px;'>\n",
        "        <b>One interface. Any environment. Zero setup.</b>\n",
        "    </p>\n",
        "</div>\n",
        "\n",
        "---\n",
        "\n",
        "## What You'll Build\n",
        "\n",
        "In this tutorial, you'll:\n",
        "1. \ud83d\udd0c **Explore OpenEnv** - Connect to BlackJack, see how the spec works\n",
        "2. \ud83c\udfb2 **Benchmark policies** - Test random vs heuristic strategies\n",
        "3. \ud83e\udde0 **Learn about GRPO** - Brief intro to the training algorithm\n",
        "4. \u26a1 **Train with Forge** - Use PyTorch's agentic RL library\n",
        "5. \ud83d\udcca **Compare results** - Measure improvement\n",
        "6. \ud83d\udd04 **Switch environments** - Show how to train on different games\n",
        "\n",
        "**This uses production code.** Same implementation as `apps/grpo/blackjack_main_fixed.py`.\n",
        "\n",
        "---\n",
        "\n",
        "### \ud83d\udcda Resources\n",
        "- \ud83d\udce6 [OpenEnv GitHub](https://github.com/meta-pytorch/OpenEnv) - Universal RL environment spec\n",
        "- \ud83d\udcc4 [GRPO Paper (arXiv:2402.03300)](https://arxiv.org/abs/2402.03300) - Group Relative Policy Optimization\n",
        "- \ud83d\udd27 [Forge GitHub](https://github.com/meta-pytorch/torchforge) - PyTorch-native agentic RL library\n",
        "- \ud83d\udcd6 [Forge Docs](https://meta-pytorch.org/torchforge/) - Full documentation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## \ud83d\udd0c Part 1: Exploring OpenEnv\n\nLet's connect to a BlackJack environment and explore the OpenEnv spec.\n\n### Start the Server\n\n<div style='background: #fff3cd; padding: 15px; border-radius: 8px; border-left: 5px solid #ffc107; margin: 20px 0;'>\n    <b>\u26a0\ufe0f Note:</b> Start the OpenEnv server in a separate terminal:\n    <pre style='margin-top: 10px; background: white; padding: 10px; border-radius: 5px;'>\n# Set your OpenEnv path\nexport OPENENV_PATH=\"/path/to/OpenEnv/src\"\nexport PYTHONPATH=\"${OPENENV_PATH}:${PYTHONPATH}\"\n\n# Start BlackJack server\nOPENSPIEL_GAME=blackjack python -m envs.openspiel_env.server.app --port 8004</pre>\n</div>"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Environment setup for Jupyter\n",
        "import sys\n",
        "import os\n",
        "\n",
        "# Fix for Monarch/Torchstore Rust bindings in Jupyter\n",
        "conda_prefix = os.environ.get('CONDA_PREFIX', sys.prefix)\n",
        "lib_path = f\"{conda_prefix}/lib\"\n",
        "\n",
        "if 'LD_LIBRARY_PATH' in os.environ:\n",
        "    if lib_path not in os.environ['LD_LIBRARY_PATH']:\n",
        "        os.environ['LD_LIBRARY_PATH'] = f\"{lib_path}:{os.environ['LD_LIBRARY_PATH']}\"\n",
        "else:\n",
        "    os.environ['LD_LIBRARY_PATH'] = lib_path\n",
        "\n",
        "print(\"\u2705 Environment configured\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Connect to OpenEnv\n",
        "\n",
        "Let's connect to the BlackJack environment and explore its interface."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import sys",
        "import os",
        "from pathlib import Path",
        "",
        "# Add OpenEnv to path (update this to your OpenEnv installation)",
        "openenv_path = os.environ.get('OPENENV_PATH', '/path/to/OpenEnv/src')",
        "if openenv_path not in sys.path:",
        "    sys.path.insert(0, openenv_path)",
        "",
        "from envs.openspiel_env import OpenSpielEnv, OpenSpielAction",
        "from grpo_utils import show_openenv_observation",
        "",
        "# Connect to environment",
        "env = OpenSpielEnv(base_url=\"http://localhost:8004\")",
        "",
        "print(\"\ud83c\udfb0 Connected to BlackJack environment\")",
        "print(\"\\n\ud83d\udccd Resetting environment...\\n\")",
        "",
        "# Reset and observe",
        "result = env.reset()",
        "show_openenv_observation(result.observation)",
        "",
        "env.close()",
        "print(\"\\n\u2705 OpenEnv interface exploration complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### What Just Happened?\n",
        "\n",
        "You just saw the **OpenEnv spec** in action:\n",
        "\n",
        "```python\n",
        "# Universal interface - works for ANY environment\n",
        "result = env.reset()              # Start episode\n",
        "result = env.step(action)         # Take action\n",
        "state = env.state()               # Get environment state\n",
        "env.close()                       # Cleanup\n",
        "```\n",
        "\n",
        "**Key observations:**\n",
        "- `legal_actions`: What actions the agent can take\n",
        "- `info_state`: Numeric observation vector\n",
        "- `game_phase`: Current phase of the game\n",
        "- `reward`: Outcome (+1 win, -1 loss, 0 push)\n",
        "\n",
        "This same interface works for **70+ different environments**. Change the server, everything else stays the same!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## \ud83c\udfb2 Part 2: Benchmarking Baseline Policies\n",
        "\n",
        "Before training an LLM, let's see how simple policies perform."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from grpo_utils import play_random_policy",
        "",
        "print(\"\ud83c\udfb2 Running random policy baseline...\\n\")",
        "",
        "# Play 100 games with random actions",
        "stats = play_random_policy(\"http://localhost:8004\", num_games=100)",
        "",
        "print(\"\\n\ud83d\udcca Random Policy Results:\")",
        "print(f\"   Games played: {stats['total_games']}\")",
        "print(f\"   Wins: {stats['wins']}\")",
        "print(f\"   Losses: {stats['losses']}\")",
        "print(f\"   Pushes: {stats['pushes']}\")",
        "print(f\"   Win rate: {stats['win_rate']:.1%}\")",
        "print(\"\\n\ud83d\udcdd Note: Optimal BlackJack strategy achieves ~43% win rate\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### The Challenge\n",
        "\n",
        "Random policy performs poorly (~30-35% win rate).\n",
        "\n",
        "**Can we train an LLM to do better?**\n",
        "\n",
        "That's where **GRPO** comes in."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## \ud83e\udde0 Part 3: Understanding Reinforcement Learning & GRPO\n\n<div style='background: linear-gradient(135deg, #e66465 0%, #9198e5 100%); padding: 25px; border-radius: 10px; color: white; margin: 20px 0; border: 3px solid #fff;'>\n    <h3 style='margin-top: 0;'>\ud83d\udcda Section Inspired by Unsloth</h3>\n    <p style='font-size: 16px; line-height: 1.8;'>\n        This section is heavily inspired by the excellent <a href='https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide' style='color: #fff; text-decoration: underline;'><b>Unsloth RL Guide</b></a>.\n        <br><br>\n        Unsloth has done an amazing job making RL accessible and intuitive. We highly recommend reading their full guide for deeper insights and practical tips!\n        <br><br>\n        \ud83d\ude4f <b>Big thanks to the Unsloth team</b> for their educational approach to RL.\n    </p>\n</div>\n\n---\n\n### What is Reinforcement Learning?\n\n<div style='background: #f8f9fa; padding: 20px; border-radius: 10px; border-left: 5px solid #6c757d; margin: 20px 0;'>\n    <h4 style='margin-top: 0;'>The Core Idea (It's Simpler Than You Think!)</h4>\n    <p style='font-size: 16px; line-height: 1.8;'>\n        The goal of RL is extremely simple:\n    </p>\n    <ul style='font-size: 16px; line-height: 1.8;'>\n        <li>\u2705 <b>Increase the chance of seeing \"good\" outcomes</b></li>\n        <li>\u274c <b>Decrease the chance of seeing \"bad\" outcomes</b></li>\n    </ul>\n    <p style='font-size: 16px; margin-top: 10px;'>\n        That's it! Everything else is just details about what \"good\" and \"bad\" mean, and how to increase/decrease their probabilities.\n    </p>\n</div>\n\n#### A Simple Example: Learning \"2 + 2 = ?\"\n\nImagine an untrained language model trying to answer \"What is 2+2?\". It might output:\n\n```\n0, cat, -10, 1928, 3, A, B, 122, 17, 182, 172, A, C, BAHS, %$, #, 9, -192, 12.31, ...\n```\n\nThen suddenly: **4** \u2713\n\nThe reward signals would be:\n```\n0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ... then 1\n```\n\n**This is the key insight:** By patience (or \"luck\"), if the correct answer has *any* non-zero probability, RL will eventually find it. The trick is:\n1. While waiting, we learn from **bad answers** \u2192 tell model \"don't do this\"\n2. When we find **good answers** \u2192 tell model \"do more of this\"\n\nThis is why I like to call it **\"Patience Is All You Need\"** for RL.\n\n---\n\n### From PPO to GRPO: The Evolution\n\n<div style='background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 25px; border-radius: 10px; color: white; margin: 20px 0;'>\n    <h4 style='margin-top: 0;'>\ud83d\udcdc The Algorithm Evolution</h4>\n    \n<table style='width: 100%; color: white; margin-top: 15px;'>\n<tr>\n    <td style='padding: 8px; border-bottom: 1px solid rgba(255,255,255,0.3);'><b>RLHF + PPO</b> (OpenAI ChatGPT)</td>\n    <td style='padding: 8px; border-bottom: 1px solid rgba(255,255,255,0.3);'>Needed 3 models: Policy, Reference, Value Model</td>\n</tr>\n<tr>\n    <td style='padding: 8px;'><b>GRPO</b> (DeepSeek R1)</td>\n    <td style='padding: 8px;'>Only needs 2 models: Policy + Reference<br>\u2192 <b>Much more efficient!</b></td>\n</tr>\n</table>\n</div>\n\n**What GRPO removes:**\n- \u274c **Value Model** \u2192 Replaced with group statistics\n- \u274c **Reward Model** \u2192 Replaced with simple reward functions\n\n**Why this matters:**\n- \ud83d\udcbe Less memory usage\n- \u26a1 Faster training\n- \ud83c\udfaf Easier to implement\n\n---\n\n### GRPO: Group Relative Policy Optimization\n\n<div style='background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); padding: 25px; border-radius: 10px; color: white; margin: 20px 0;'>\n    <h4 style='margin-top: 0;'>Why \"Group Relative\"?</h4>\n    <p style='font-size: 16px; line-height: 1.8;'>\n        Instead of training a separate Value Model to estimate \"how good is this state?\", \n        GRPO uses a clever trick: <b>sample the model multiple times</b> and compare answers within the group.\n    </p>\n</div>\n\n**Example: Training on \"What is 2+2?\"**\n\n1. **Generate multiple responses** (e.g., 4 samples):\n   - Response 1: \"4\" \u2192 reward = +1 (correct!)\n   - Response 2: \"3\" \u2192 reward = 0 (close, but wrong)\n   - Response 3: \"D\" \u2192 reward = -1 (nonsense)\n   - Response 4: \"C\" \u2192 reward = -1 (nonsense)\n\n2. **Calculate group statistics:**\n   - Mean reward: (-1 + -1 + 0 + 1) / 4 = -0.25\n   - Standard deviation: ~0.83\n\n3. **Compute advantages** (Z-score normalization):\n   - Response 1: +1.5 (much better than average!)\n   - Response 2: +0.3 (slightly better)\n   - Response 3: -0.9 (worse than average)\n   - Response 4: -0.9 (worse than average)\n\n4. **Update model:**\n   - Increase probability of generating \"4\"\n   - Slightly increase \"3\" (it's closer than nonsense)\n   - Decrease probability of generating \"D\" and \"C\"\n\nThis is **group-relative** because we're comparing within the group, not to an absolute baseline!\n\n---\n\n### Reward Functions: The Secret Sauce\n\nReward functions tell the model what's \"good\" and what's \"bad\". They can be simple or complex:\n\n**For BlackJack (what we're using):**\n```python\ndef evaluate_response(prompt, response, game_reward):\n    reward = float(game_reward)  # +1 (win), -1 (loss), 0 (push)\n    \n    # Reward shaping: Scale up wins\n    if game_reward > 0:\n        reward = 2.0  # Wins are more valuable\n    elif game_reward == 0:\n        reward = 0.5  # Pushes better than losses\n    \n    return reward\n```\n\n**For Math Problems:**\n- If answer is a number: +1\n- If answer matches ground truth: +3\n- If no number detected: -1\n- **Total reward:** Sum of all criteria\n\n**For Email Automation:**\n- Contains required keyword: +1\n- Matches ideal response: +1\n- Too long: -1\n- Includes recipient name: +1\n- Has signature block: +1\n\nThe key is: **Reward functions must be verifiable**. You can't subjectively judge \"is this creative?\" but you can verify \"is this answer correct?\"\n\n---\n\n### The Training Process (Simplified)\n\n```\n1. Play game \u2192 Get action \"HIT\" or \"STAND\"\n   \u2193\n2. Game ends \u2192 Observe reward (+1 win, -1 loss, 0 push)\n   \u2193\n3. Repeat 4-8 times for the same question (group)\n   \u2193\n4. Calculate group statistics (mean, std)\n   \u2193\n5. Compute advantages (which answers were better/worse than average?)\n   \u2193\n6. Update model: increase good action probability, decrease bad\n   \u2193\n7. Repeat thousands of times \u2192 Model learns strategy!\n```\n\n**Key insight:** Over time, the model learns not just \"what to do\" but also *why* (the reasoning process). This is how DeepSeek R1 developed its famous `<think>` tokens!\n\n---\n\n### Forge: PyTorch-Native Agentic RL Infrastructure\n\n<div style='background: linear-gradient(135deg, #20c997 0%, #17a2b8 100%); padding: 20px; border-radius: 10px; color: white; margin: 20px 0;'>\n    <h4 style='margin-top: 0;'>What is Forge?</h4>\n    <p style='font-size: 16px; line-height: 1.6;'>\n        <b>Forge</b> is PyTorch's official library for training agentic RL models. It handles all the distributed systems complexity so you can focus on algorithms.\n    </p>\n    <ul style='font-size: 15px; line-height: 1.7;'>\n        <li><b>Generator (vLLM):</b> Fast LLM inference with automatic batching</li>\n        <li><b>RLTrainer:</b> Distributed training with FSDP across GPUs</li>\n        <li><b>ReplayBuffer:</b> Stores episodes for off-policy learning</li>\n        <li><b>ReferenceModel:</b> Keeps original model for KL penalty</li>\n        <li><b>Torchstore:</b> Distributed weight management across replicas</li>\n    </ul>\n</div>\n\n**Resources:**\n- \ud83d\udd27 [GitHub](https://github.com/meta-pytorch/torchforge) - Source code\n- \ud83d\udcd6 [Documentation](https://meta-pytorch.org/torchforge/) - Full docs\n- \ud83d\udcc4 [GRPO Paper](https://arxiv.org/abs/2402.03300) - Original research\n\n**In this tutorial:** We abstract all of Forge's complexity. You just call:\n```python\ntrainer = await setup_forge_training(\"config.yaml\")\nawait trainer.run(steps=100)\n```\n\nEverything else happens automatically! \ud83d\ude80"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## \ud83c\udfd7\ufe0f Part 4: Training with GRPO\n",
        "\n",
        "Now let's train a Qwen 1.5B model to play BlackJack using production GRPO code.\n",
        "\n",
        "### Architecture Overview\n",
        "\n",
        "```\n",
        "\u250f\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2513\n",
        "\u2503              YOUR TRAINING LOOP                    \u2503\n",
        "\u2523\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u252b\n",
        "\u2503                                                    \u2503\n",
        "\u2503  Rollouts Loop          Training Loop             \u2503\n",
        "\u2503  \u2022 Play games           \u2022 Sample batch            \u2503\n",
        "\u2503  \u2022 Collect episodes     \u2022 Compute loss            \u2503\n",
        "\u2503  \u2022 Compute advantages   \u2022 Update weights          \u2503\n",
        "\u2503  \u2022 Add to buffer        \u2022 Push to replicas        \u2503\n",
        "\u2503                                                    \u2503\n",
        "\u2517\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u252f\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u252f\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u251b\n",
        "           \u2502                         \u2502\n",
        "      HTTP \u2502                         \u2502 RPC\n",
        "           \u2502                         \u2502\n",
        "           \u2193                         \u2193\n",
        "   \u250f\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2513          \u250f\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2513\n",
        "   \u2503   OpenEnv   \u2503          \u2503    Forge     \u2503\n",
        "   \u2503   Server    \u2503          \u2503   Services   \u2503\n",
        "   \u2517\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u251b          \u2517\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u251b\n",
        "```\n",
        "\n",
        "**Two concurrent loops:**\n",
        "1. **Rollouts:** Play games via OpenEnv \u2192 collect episodes\n",
        "2. **Training:** Sample from buffer \u2192 update policy with GRPO\n",
        "\n",
        "They run in parallel for maximum efficiency!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Setup and Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from grpo_utils import setup_forge_trainingprint(\"\ud83c\udfd7\ufe0f Initializing Forge infrastructure...\\n\")print(\"This will:\")print(\"  \u2022 Load the Qwen 1.5B model\")print(\"  \u2022 Initialize vLLM inference servers\")print(\"  \u2022 Setup distributed training (TorchTitan)\")print(\"  \u2022 Create replay buffer and reference model\")print(\"\\n\u23f3 This may take 1-2 minutes...\\n\")# Initialize everything with one function calltrainer = await setup_forge_training(\"blackjack.yaml\")print(\"\\n\u2705 Ready to train!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Run Training\n",
        "\n",
        "Now we train for 100 steps. This is a shortened demo - production training uses 1000+ steps."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "print(\"\ud83d\ude80 Starting GRPO training!\\n\")\n",
        "print(\"Watch the logs to see:\")\n",
        "print(\"  \u2022 Games being played (with actions and outcomes)\")\n",
        "print(\"  \u2022 Win rate improving over time\")\n",
        "print(\"  \u2022 Training steps updating the policy\")\n",
        "print(\"\\n\" + \"=\"*60 + \"\\n\")\n",
        "\n",
        "# Run training (this is the production training loop!)\n",
        "results = await trainer.run(steps=100)\n",
        "\n",
        "print(\"\\n\" + \"=\"*60)\n",
        "print(\"\\n\ud83c\udf89 Training complete!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Cleanup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Shutdown Forge services\n",
        "await trainer.shutdown()\n",
        "print(\"\u2705 Shutdown complete\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## \ud83d\udd04 Part 5: The Power of OpenEnv - Switching Environments\n",
        "\n",
        "Here's the magic: **The same code works for ANY OpenEnv environment.**\n",
        "\n",
        "### Switch to Tic-Tac-Toe\n",
        "\n",
        "Just change the server:\n",
        "\n",
        "```bash\n",
        "# Terminal:\n",
        "OPENSPIEL_GAME=tic_tac_toe python -m envs.openspiel_env.server.app --port 8005\n",
        "```\n",
        "\n",
        "Update config:\n",
        "```python\n",
        "cfg.blackjack_env.server_url = \"http://localhost:8005\"\n",
        "```\n",
        "\n",
        "**Everything else stays identical.** Same GRPO code, same Forge infrastructure.\n",
        "\n",
        "---\n",
        "\n",
        "### Switch to Chess\n",
        "\n",
        "```bash\n",
        "OPENSPIEL_GAME=chess python -m envs.openspiel_env.server.app --port 8006\n",
        "```\n",
        "\n",
        "Update model and config for longer sequences, done!\n",
        "\n",
        "---\n",
        "\n",
        "### Switch to Atari\n",
        "\n",
        "```bash\n",
        "# Different OpenEnv backend\n",
        "python -m envs.atari_env.server.app --game pong --port 8007\n",
        "```\n",
        "\n",
        "Modify prompt formatting for vision inputs, same training loop!\n",
        "\n",
        "---\n",
        "\n",
        "<div style='background: #d1ecf1; padding: 20px; border-radius: 10px; border-left: 5px solid #0c5460; margin: 20px 0;'>\n",
        "    <h3 style='color: #0c5460; margin-top: 0;'>\ud83d\udca1 The Key Insight</h3>\n",
        "    <p style='color: #0c5460; font-size: 16px;'>\n",
        "        <b>OpenEnv is a spec, not a game engine.</b><br><br>\n",
        "        Once you have a training loop that talks to OpenEnv, you can train on ANY environment that implements the spec.\n",
        "        <br><br>\n",
        "        Change one environment variable \u2192 train on 70+ different environments.\n",
        "    </p>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## \ud83d\ude80 Next Steps\n",
        "\n",
        "### 1. Scale Up Training\n",
        "\n",
        "Edit `apps/grpo/blackjack.yaml`:\n",
        "\n",
        "```yaml\n",
        "trainer:\n",
        "  training:\n",
        "    steps: 1000          # More training steps\n",
        "\n",
        "group_size: 8            # More games per rollout\n",
        "rollout_threads: 4       # Parallel rollout collection\n",
        "```\n",
        "\n",
        "Run from command line for serious training:\n",
        "\n",
        "```bash\n",
        "python -m apps.grpo.blackjack_main_fixed --config apps/grpo/blackjack.yaml\n",
        "```\n",
        "\n",
        "---\n",
        "\n",
        "### 2. Explore Other Environments\n",
        "\n",
        "Try different OpenSpiel games:\n",
        "- `OPENSPIEL_GAME=tic_tac_toe`\n",
        "- `OPENSPIEL_GAME=connect_four`\n",
        "- `OPENSPIEL_GAME=go`\n",
        "\n",
        "Explore other OpenEnv backends:\n",
        "- Atari environments\n",
        "- FinRL trading simulations\n",
        "- Custom environments\n",
        "\n",
        "---\n",
        "\n",
        "### 3. Customize the Training\n",
        "\n",
        "All the code is in `apps/grpo/grpo_utils.py`:\n",
        "- Modify reward shaping in `BlackJackReward.evaluate_response()`\n",
        "- Adjust advantage computation in `ComputeAdvantages.compute()`\n",
        "- Tweak GRPO loss hyperparameters (beta, KL penalty)\n",
        "- Change prompt formatting in `format_prompt()`\n",
        "\n",
        "---\n",
        "\n",
        "## \ud83d\udcda Resources\n",
        "\n",
        "### OpenEnv\n",
        "- \ud83d\udce6 [GitHub](https://github.com/meta-pytorch/OpenEnv) - Source code and examples\n",
        "- \ud83d\udcd6 [Spec Documentation](https://github.com/meta-pytorch/OpenEnv#spec) - Full API reference\n",
        "\n",
        "### GRPO\n",
        "- \ud83d\udcc4 [Paper (arXiv:2402.03300)](https://arxiv.org/abs/2402.03300) - Original publication\n",
        "- \ud83d\udd2c [Blog Post](https://ai.meta.com/blog/grpo/) - High-level explanation\n",
        "\n",
        "### Forge\n",
        "- \ud83d\udd27 [GitHub](https://github.com/meta-pytorch/torchforge) - PyTorch-native agentic RL\n",
        "- \ud83d\udcd6 [Docs](https://meta-pytorch.org/torchforge/) - Full documentation\n",
        "- \ud83d\udcac [Discussions](https://github.com/meta-pytorch/torchforge/discussions) - Community support\n",
        "\n",
        "---\n",
        "\n",
        "## \ud83c\udf93 Key Takeaways\n",
        "\n",
        "<div style='background: #d4edda; padding: 25px; border-radius: 10px; border-left: 5px solid #28a745; margin: 20px 0;'>\n",
        "    <h3 style='color: #155724; margin-top: 0;'>What You Learned</h3>\n",
        "    <ol style='color: #155724; font-size: 16px; line-height: 1.8;'>\n",
        "        <li><b>OpenEnv is a universal spec</b> for RL environments - not just games, ANY interactive environment.</li>\n",
        "        <li><b>One training loop works everywhere</b> - switch environments by changing a URL.</li>\n",
        "        <li><b>Forge abstracts distributed RL complexity</b> - focus on algorithms, not infrastructure.</li>\n",
        "        <li><b>GRPO enables stable LLM training</b> - group-relative advantages + KL penalties work.</li>\n",
        "        <li><b>Production code is accessible</b> - this notebook uses the same code as large-scale training.</li>\n",
        "    </ol>\n",
        "</div>\n",
        "\n",
        "---\n",
        "\n",
        "<div style='background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 30px; border-radius: 10px; color: white; margin: 30px 0; text-align: center;'>\n",
        "    <h2 style='margin-top: 0;'>\ud83c\udf89 Congratulations!</h2>\n",
        "    <p style='font-size: 18px; line-height: 1.8;'>\n",
        "        You just trained an LLM using production GRPO code.<br>\n",
        "        You explored OpenEnv as a universal RL interface.<br>\n",
        "        You saw how Forge abstracts distributed training complexity.\n",
        "    </p>\n",
        "    <p style='font-size: 20px; margin-top: 20px;'>\n",
        "        <b>Now go train agents in ANY environment! \ud83d\ude80</b>\n",
        "    </p>\n",
        "</div>"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}